package main

import (
	"bufio"
	"fmt"
	"os"
	"runtime"
	"strings"
)

func getHostsFilePath() string {
	switch runtime.GOOS {
	case "linux":
		return "/etc/hosts"
	case "windows":
		return `C:\Windows\System32\drivers\etc\hosts`
	default:
		return ""
	}
}

func readHostsFile(filePath string) ([]string, error) {
	var lines []string
	file, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if len(line) > 0 && !strings.HasPrefix(line, "#") {
			lines = append(lines, line)
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return lines, nil
}

func addOrUpdateHostEntry(hostEntry HostEntry) error {
	hostsFilePath := getHostsFilePath()
	if hostsFilePath == "" {
		return fmt.Errorf("unknown operating system")
	}

	lines, err := readHostsFile(hostsFilePath)
	if err != nil {
		return err
	}

	for _, line := range lines {
		if strings.Fields(line)[1] == hostEntry.Host {
			err := deleteHostsEntry(hostEntry)
			if err != nil {
				return err
			}
		}
	}

	file, err := os.OpenFile(hostsFilePath, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer file.Close()

	entry := hostEntry.IP + " " + hostEntry.Host + "\n"
	_, err = file.WriteString(entry)
	if err != nil {
		return err
	}

	return nil
}

func deleteHostsEntry(host HostEntry) error {

	hostsFilePath := getHostsFilePath()
	if hostsFilePath == "" {
		return fmt.Errorf("unknown operating system")
	}

	lines, err := readHostsFile(hostsFilePath)
	if err != nil {
		return err
	}

	// 遍历每一行，查找匹配的域名并删除
	var updatedLines []string
	for _, line := range lines {
		parts := strings.Fields(line)
		if len(parts) > 1 && parts[1] == host.Host {
			continue // 跳过匹配的行
		}
		updatedLines = append(updatedLines, line+"\n")
	}

	file, err := os.OpenFile(hostsFilePath, os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer file.Close()

	// 将更新后的内容写回hosts文件
	_, err = file.WriteString(strings.Join(updatedLines, ""))
	if err != nil {
		return err
	}

	return nil
}

type HostEntry struct {
	IP   string
	Host string
}
